package core.matrix;

import core.iterfaces.Matrix;
import core.iterfaces.Vec;
import core.vec.DenseVec;
import core.vec.SparseVec;
import exception.NotSupportException;

/**
 * @author wzd
 * 带状矩阵，具有下带宽p和上带宽q，且p和q远小于n.
 * 这样的矩阵可以按照
 * a_{ij}=A.band(i-j+q+1,j)
 * 存储在一个（p+q+1）xn的二维数组A。band
 */

public class BandMatrix implements Matrix {
    /**
     * 下带宽p
     */
    private Integer p;
    /**
     * 上带宽q
     */
    private Integer q;
    private Double[][] band;
    private Integer row;
    private Integer col;

    public BandMatrix(Integer row,Integer col,Integer p,Integer q){
        this.row = row;
        this.col = col;
        this.p = p;
        this.q = q;
        band = new Double[p+q+1][col];
        /**
         * 初始化元素为0
         */
        for (int i = 0; i < p+q+1; i++) {
            for (int j = 0; j < col; j++) {
                band[i][j] = 0.0;
            }
        }
    }

    /**
     * 构造函数，使用d
     * @param row
     * @param col
     * @param p
     * @param q
     * @param data
     */
    public BandMatrix(Integer row,Integer col,Integer p,Integer q,Double[][] data){
        this.row = row;
        this.col = col;
        this.p = p;
        this.q = q;
        band = new Double[p+q+1][col];
        /**
         * 初始化元素
         * A.band[i][j]
         */
        for (int i = 0; i < p+q+1; i++) {
            for (int j = 0; j < col; j++) {
                int k = i + j - q;
                if(k>=0&&k<row) {
                    band[i][j] = data[i + j - q][j];
                }
                else{
                    band[i][j] = 0.0;
                }
            }
        }
    }

    @Override
    public Matrix zero() {
        return new BandMatrix(this.row,this.col,this.p,this.q);
    }

    @Override
    public Matrix zero(int rowLen, int colLen) {
        return new BandMatrix(rowLen,colLen,this.p,this.q);
    }

    @Override
    public double get(int i, int j) {
        /**
         * 上下带宽之外，全部为0
         */
        if(i-j>p||j-i>q){
            return 0.0;
        }
        /**
         * 否则，按照aij=band[i-j+q][j]存储
         */
        return band[i-j+q][j];
    }

    @Override
    public void set(int i, int j, double value) {
        if(i-j>p||j-i>q){
        }
        else{
            band[i-j+q][j] = value;
        }
    }

    @Override
    public Vec getRow(int i) throws NotSupportException {
        /**
         * 返回的结果，必然是稀疏向量
         */
        Vec result = new SparseVec(col);
        /**
         * 第i行j从i-p和0较大的一个位置开始，直到i+q和n较小的一个
         */
        Integer from = Math.max(0,i-p);
        Integer to = Math.min(i+q,col-1);
        for (int j = from; j < to+1; j++) {
            result.set(j,get(i,j));
        }
        return result;
    }

    @Override
    public Vec getCol(int j) {
        /**
         * 返回的结果，必然是稀疏向量
         */
        Vec result = new SparseVec(col);
        /**
         * 从第j列i从max(j-q,0)到min(j+p,m)
         */
        Integer from = Math.max(j-q,0);
        Integer to = Math.min(j+p,row-1);
        for (int i = from; i < to+1; i++) {
            result.set(i,get(i,j));
        }
        return result;
    }

    @Override
    public Vec getSubRow(int i, int begin, int end) throws NotSupportException {
        if(i<0||i>row-1){
            throw new IndexOutOfBoundsException("传入的参数i:"+i+" 不符合矩阵行数范围");
        }
        if(begin<0||begin>col-1){
            throw new IndexOutOfBoundsException("传入的参数begin:"+begin+" 不符合矩阵列数范围");
        }
        if(end<0||end>col){
            throw new IndexOutOfBoundsException("传入的参数end:"+end+" 不符合矩阵列数范围");
        }
        Integer length = end - begin;
        Vec result = new SparseVec(length);
        /**
         * 算出当前行最左和最有的下标记
         */
        Integer from = Math.max(0,i-p);
        Integer to = Math.min(i+q,col-1);
        if(begin<from){begin = from;}
        if(end>to+1){
            end = to+1;
        }
        for (int j = begin; j < end; j++) {
            result.set(j,get(i,j));
        }
        return result;
    }

    @Override
    public Vec getSubCol(int j, int begin, int end) throws NotSupportException {
        if(j<0||j>col-1){
            throw new IndexOutOfBoundsException("传入的参数j:"+j+" 不符合矩阵列数范围");
        }
        if(begin<0||begin>row-1){
            throw new IndexOutOfBoundsException("传入的参数begin:"+begin+" 不符合矩阵行数范围");
        }
        if(end<0||begin>row){
            throw new IndexOutOfBoundsException("传入的参数end:"+end+" 不符合矩阵行数范围");
        }
        Integer length = end - begin;
        Vec result = new SparseVec(length);
        Integer from = Math.max(j-q,0);
        Integer to = Math.min(j+p,row-1);
        if(begin<from){begin = from;}
        if(end>to+1){
            end = to+1;
        }
        for (int i = begin; i < end; i++) {
            result.set(i,get(i,j));
        }
        return result;
    }

    @Override
    public int getRowLen() {
        return row;
    }

    @Override
    public int getColLen() {
        return col;
    }

    @Override
    public boolean isDense() {
        return false;
    }

    @Override
    public boolean isPositiveDefinite() {
        return false;
    }

    @Override
    public double norm2() {
        return 0;
    }

    @Override
    public Matrix toSparse() throws NotSupportException {
        SparseMatrix result = new SparseMatrix(this.row,this.col);
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                result.set(i,j, get(i,j));
            }
        }
        return result;
    }

    @Override
    public Matrix toDense() throws NotSupportException {
        DenseMatrix result = new DenseMatrix(this.row,this.col);
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                result.set(i,j, get(i,j));
            }
        }
        return result;
    }

    /**
     * 返回对角线元素
     * @param k 对角线的位置 k=0为对角线 k>0 为上对角线 k<0为下对角线
     * @return
     */
    @Override
    public Vec getKDiag(int k) {
        int len = diagLength(k);
            Double[] data = band[q-k];
            Vec result = new DenseVec(len);
            for(int i=0;i<len;i++){
                if(k>0&&k<=q) {
                    result.set(i, data[i + k]);
                }
                else{
                    result.set(i,data[i]);
                }
            }
            return result;
    }

    /**
     * 带状矩阵只涉及带状部分的反转，
     * @return
     */
    @Override
    public Matrix transform() {
        Matrix result = new BandMatrix(col,row,q,p);
        for (int i = 0; i < p+q+1; i++) {
            for (int j = 0; j < col; j++) {
                int k = i + j - q;
                if(k>=0&&k<row) {
                    result.set(j,i + j - q,band[i][j]);
                }
            }
        }
        return result;
    }

    @Override
    public Matrix muti(double a, boolean isCovered) {
        if(isCovered){
            for (int i = 0; i < p+q+1; i++) {
                for (int j = 0; j < col; j++) {
                    band[i][j]=a*band[i][j];
                }
            }
            return this;
        }
        else{
            Matrix result = zero();
            for (int i = 0; i < p+q+1; i++) {
                for (int j = 0; j < col; j++) {
                    int k = i + j - q;
                    if(k>=0&&k<row) {
                        result.set(j,i + j - q,band[i][j]);
                    }
                }
            }
            return result;
        }
    }

    /**
     * 使用A.band相乘，使用向量元素是矩阵列线性组合的系数这个性质
     * @param v 与矩阵相乘的向量
     * @return
     * @throws NotSupportException
     */
    @Override
    public Vec muti(Vec v) throws NotSupportException {
        Vec result = new DenseVec(row);
        for (int j = 0; j < col; j++) {
            int top = Math.max(0,j-q);
            int bot = Math.min(row,j+p+1);
            int bandTop = Math.max(0,q-j);
            int bandBot = bandTop+bot-top;
            for (int k=top;k<bot;k++){
                result.set(k,v.get(j)*band[bandTop+k-top][j]+result.get(k));
            }
        }
        return result;
    }

    @Override
    public Matrix muti(Matrix other) throws NotSupportException {
        if(this.col!= other.getRowLen()){
            throw new NotSupportException("矩阵不相容，不能相乘");
        }
        Matrix result = new DenseMatrix(row,other.getColLen());
        for(int j=0;j<other.getColLen();j++){
           Vec colJ = muti(other.getCol(j));
            /**
             * 设置colJ
             */
            for (int i = 0; i < row; i++) {
                result.set(i,j,colJ.get(i));
            }
        }
        return result;
    }



    @Override
    public boolean isLowerTri() {
        return q==0;
    }

    @Override
    public boolean isUpperTri() {
        return p==0;
    }

    @Override
    public boolean isHessenberg() {
        return p<=1;
    }

    @Override
    public boolean isSymmetric() {
        if(p!=q){
            return false;
        }
        for (int i = 0; i < q; i++) {
            for (int j = 0; j < col; j++) {
                int k = i + j - q;
                if(k>=0&&k<row) {
                    if(band[i][j] != get(j,i + j - q)){
                        return false;
                    }
                }
            }
        }
        return true;
    }

    @Override
    public boolean isTriDiag() {
        return p==1&&q==1;
    }


    @Override
    public double fnorm() {
        Double result = 0.0;
        for (int i = 0; i < p+q+1; i++) {
            for (int j = 0; j < col; j++) {
                result+=Math.pow(band[i][j],2);
            }
        }
        return result;
    }
}
